﻿//////////////////////////////////////////////////////////////////////
// Constraint programming - Finite Domains
// 
// Ref: http://overtond.blogspot.com/2008/07/pre.html
//////////////////////////////////////////////////////////////////////

module CPFD.Core

open Monad

let cp = StateListMonad()

type FDVar = { id : int }
and VarInfo = { values : Set<int>; delayedConstraints : FDState -> List<unit*FDState> }
and FDState = { varMap : Map<FDVar, VarInfo>; nextId : int }

let runFD fd = cp.Execute(fd, { varMap = Map.empty; nextId = 0 })

let newVar domain = 
    let nextVar = cp {
        let! s = cp.Get()
        let vid = s.nextId
        do! cp.Set({ s with nextId = vid + 1})
        return { id = vid }
    }
    let isOneOf x domain = cp {
        return! cp.Update(fun s ->
            let vm = s.varMap
            let vi = { values = Set.ofList domain; delayedConstraints = cp.Return(()) }
            { s with varMap = Map.add x vi vm }
        )
    }    
    cp {
        let! v = nextVar
        do! isOneOf v domain
        return v
    }

let rec newVars n domain = cp.replicateM n (newVar domain)

let lookup x = cp {
        let! s = cp.Get()
        return s.varMap.[x].values
    }

let update x i = cp {
        let! s = cp.Get()
        let vm = s.varMap
        let vi = vm.[x]
        do! cp.Set({ s with varMap = Map.add x { vi with values = i } vm })
        do! vi.delayedConstraints
    }

let addConstraint x constr = cp {
        let! s = cp.Get()
        let vm = s.varMap
        let vi = vm.[x]
        let delayedConstraints' = cp { 
            do! constr
            do! vi.delayedConstraints 
        }
        do! cp.Set({ s with varMap = Map.add x { vi with delayedConstraints = delayedConstraints' } vm })
    }
    
let addBinaryConstraint f x y = cp {
        let constr = f x y
        do! constr
        do! addConstraint x constr
        do! addConstraint y constr
    }

// experimental
let addMultiVariableConstraint f xs = cp {
        let constr = f xs
        do! constr
        for v in xs do
            do! addConstraint v constr 
    }

let hasValue var value = cp {
        let! vals = lookup var
        do! cp.Guard(Set.contains value vals)
        let i = Set.singleton value        
        if i <> vals then 
            do! update var i
    }

let same = 
    let sameConstraint x y = cp {
        let! xv = lookup x
        let! yv = lookup y
        let i = Set.intersect xv yv
        do! cp.Guard(Set.isEmpty i |> not)         
        if i <> xv then 
            do! update x i
        if i <> yv then 
            do! update y i
    }
    addBinaryConstraint sameConstraint

let compareBy f =    
    let sameByConstraint x y = cp {
        let! xv = lookup x
        let! yv = lookup y        
        let validElems = seq {
            for v1 in xv do
                for v2 in yv do
                    if f v1 v2 then
                        yield (v1,v2)
        }
        let xv' = Set.intersect xv (Seq.map fst validElems |> Set.ofSeq)
        let yv' = Set.intersect yv (Seq.map snd validElems |> Set.ofSeq)
        do! cp.Guard(not(Set.isEmpty xv') && not(Set.isEmpty yv'))         
        if xv' <> xv then 
            do! update x xv'
        if yv' <> yv then 
            do! update y yv'
    }
    addBinaryConstraint sameByConstraint

let different =
    let differentConstraint x y = cp {
        let! xv = lookup x
        let! yv = lookup y
        do! cp.Guard(Set.count xv > 1 || Set.count yv > 1 || xv <> yv)
        if Set.count xv = 1 && Set.isProperSubset xv yv then
            do! update y (Set.difference yv xv)
        if Set.count yv = 1 && Set.isProperSubset yv xv then
            do! update x (Set.difference xv yv)
    }
    addBinaryConstraint differentConstraint
    
let rec allDifferent xs = 
    cp {
        match xs with
        | x::rest -> 
            do! cp.mapM_ (different x) rest
            do! allDifferent rest
        | _ -> return ()
    }    
     
let rec labelling vars =
    let label var = cp {
        let! vals = lookup var
        let! value = cp.Lift(Set.toList vals)
        do! hasValue var value
        return value
    }
    cp.mapM label vars


